[6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto#1928
[6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto#1928galagam wants to merge 2 commits into
Conversation
|
/claude review |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds a helper to detect oversized in-memory ONNX models, routes those models through temporary external-data files for TensorRT parsing, and adds a regression test that compares parsing results across modes. ChangesFile-backed ONNX parsing support
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant get_custom_layers
participant _requires_file_backed_parse
participant TempFileSave
participant TrtParser
Caller->>get_custom_layers: onnx_path (str or ModelProto)
alt onnx_path is string
get_custom_layers->>TrtParser: parse_from_file(onnx_path)
else onnx_path is ModelProto
get_custom_layers->>_requires_file_backed_parse: check ByteSize()
alt requires file-backed parse
_requires_file_backed_parse-->>get_custom_layers: True
get_custom_layers->>TempFileSave: deepcopy + save external data
get_custom_layers->>TrtParser: parse_from_file(temp_path)
else
_requires_file_backed_parse-->>get_custom_layers: False
get_custom_layers->>TrtParser: parse(SerializeToString())
end
end
TrtParser-->>get_custom_layers: parsed layers or errors
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| with tempfile.TemporaryDirectory(prefix="modelopt_trt_") as tmpdir: | ||
| model_path = os.path.join(tmpdir, "model.onnx") | ||
| # Deep-copy so externalization doesn't strip weights from the caller's model. | ||
| model_copy = copy.deepcopy(onnx_path) |
There was a problem hiding this comment.
[SUGGESTION] copy.deepcopy(onnx_path) transiently doubles peak RAM for exactly the models this path targets — an 8GB ModelProto needs ~16GB before onnx.save streams weights to disk. The deep-copy is necessary for correctness (onnx.save(save_as_external_data=True) mutates the proto in place, stripping raw_data from the caller's model), so this isn't a bug — but on the 8GB VLA trunk case cited in the PR description it may be tight on constrained hosts.
If it ever bites, an alternative that avoids the copy is to externalize in place and restore afterward (e.g. onnx.save(onnx_path, ...) then onnx.load_external_data_for_model(onnx_path, tmpdir) to re-hydrate the caller's proto), or wrap in try/finally to guarantee re-hydration on parse failure. Non-blocking; the current approach is the safest correctness-wise.
There was a problem hiding this comment.
The peak memory usage is at ReferenceRunner where we define an auxiliary graph with every intermediate tensor being a graph output. The 2x memory usage here is not a bottleneck.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Reviewed the 2 files in this PR's authoritative changed-file list (modelopt/onnx/trt_utils.py, tests/gpu/onnx/quantization/test_plugin.py). The two-dot git diff origin/main HEAD surfaced many unrelated files due to the shallow-checkout divergence between the local origin/main tip and the true merge base — those are not part of this PR and were excluded.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Assessment — low risk. Focused, well-targeted bugfix. get_custom_layers previously silently returned an empty network (0 layers/0 tensors) for in-memory ModelProtos at/above the protobuf 2 GiB limit because TensorRT's parser.parse(bytes) fails quietly while SerializeToString() succeeds. The fix routes over-limit (or un-sizeable) in-memory models through a temporary external-data file + parse_from_file(), matching the existing string-path branch and ModelOpt's file-backed policy elsewhere in this module.
Verified:
onnx.checker.MAXIMUM_PROTOBUFis a valid constant, already used with the same>= MAXIMUM_PROTOBUFidiom at trt_utils.py:382.- The
ByteSize()-overflowexceptbranch conservatively routes to the file path, consistent withload_onnx_model's handling. copy.deepcopycorrectly shields the caller's model from in-place externalization byonnx.save; both in-memory callers pass a model the caller still owns.- The regression test forces the file-backed path via monkeypatch and cross-checks results against both the in-memory fast path and the string-path baseline. Backward compatible — the small-model fast path is unchanged.
The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models — non-blocking, and the current approach is the correctness-safe choice.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1928 +/- ##
==========================================
- Coverage 77.73% 77.71% -0.02%
==========================================
Files 519 519
Lines 57886 57902 +16
==========================================
+ Hits 45000 45001 +1
- Misses 12886 12901 +15
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
gcunhase
left a comment
gcunhase
left a comment
There was a problem hiding this comment.
Do you think the comment from Claude review should be a concern?
The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models.
Otherwise, LGTM!
@gcunhase Thanks for holding me accountable... Replied in the original thread. |
galagam
force-pushed
the
dev-gagam-bug-6403893
branch
from
b52dc3f to
c8a5a30
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gpu/onnx/quantization/test_plugin.py`:
- Line 163: Move the `modelopt.onnx.trt_utils as trt_utils` import out of the
test body and into the module-level imports in `test_plugin.py`. This import is
only used to support `monkeypatch.setattr`, so it should be declared alongside
the other top-level `modelopt.onnx.trt_utils` imports rather than inside the
test. Remove the in-function import and keep the `trt_utils` reference available
for the existing monkeypatch usage.
🪄 Autofix (Beta)
❌ Autofix failed (check again to retry)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e2d67f9b-19fd-4e4f-a16a-79d54c790fd8
📒 Files selected for processing (2)
modelopt/onnx/trt_utils.pytests/gpu/onnx/quantization/test_plugin.py
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. An unexpected error occurred while generating fixes: Not Found - https://docs.github.com/rest/git/refs#get-a-reference |
…Proto get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com>
galagam
force-pushed
the
dev-gagam-bug-6403893
branch
from
d5228a7 to
55b0296
Compare
2 participants
What does this PR do?
Type of change: ? Bugfix
get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path.
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit
Summary
Bug Fixes
Tests
get_custom_layersresults for in-memory models versus file-path parsing, including a forced file-backed routing scenario.